
#include <vector>
#include "AESSymmKey.h"
#include "RSAPublicKey.h"
#include "RSAPrivateKey.h"
#include "CryptEChain.h"
#include "EncryptWrapper.h"

void EncryptWrapper::InitEncryptWrapper(const char *iekeyname, EncryptionKey *enckey, const char *iskeyname, SigningKey *sgnkey, int compress)
{
   if (iekeyname)
   {
      ekeyname = (char *)malloc(sizeof(char)*(strlen(iekeyname)+1));
      strcpy(ekeyname, iekeyname);
   } else
   {
      ekeyname = NULL;
   }
   if (enckey)
   {
      switch (enckey->GetKeyType())
      {
         case KEY_TYPE_PUBLIC_RSA:
         {
            RSAPublicKey *rsapubkey =
               new RSAPublicKey(*((RSAPublicKey *)enckey));
            enckeycopy = (RSAPublicKey *)rsapubkey;
            break;
         }
         case KEY_TYPE_SYMM_AES:
         {
            AESSymmKey *aessymmkey =
               new AESSymmKey(*((AESSymmKey *)enckey));
            enckeycopy = (AESSymmKey *)aessymmkey;
            break;
         }
         default:
         {
            errprint("unknown encryption key type");
            enckeycopy = NULL;
            break;
         }
      }
   } else
   {
      errprint("encrypt wrapper has no key: can't do encryption!");
      enckeycopy = NULL;
   }
   if (iskeyname)
   {
      skeyname = (char *)malloc(sizeof(char)*(strlen(iskeyname)+1));
      strcpy(skeyname, iskeyname);
   } else
   {
      skeyname = NULL;
   }
   if (sgnkey)
   {
      switch (sgnkey->GetKeyType())
      {
         case KEY_TYPE_PRIVATE_RSA:
         {
            RSAPrivateKey *rsaprvkey =
               new RSAPrivateKey(*((RSAPrivateKey *)sgnkey));
            sgnkeycopy = (RSAPrivateKey *)rsaprvkey;
            break;
         }
         default:
         {
            errprint("unknown signing key type");
            sgnkeycopy = NULL;
            break;
         }
      }
   } else
   {
      sgnkeycopy = NULL;
   }
   docompress = (compress != 0);
   trusted = 0;
   apr_status_t status;
   status = apr_pool_create(&mp, NULL);
   CheckAPRError(status);
}

EncryptWrapper::EncryptWrapper(const char *iekeyname, EncryptionKey *enckey)
{
   InitEncryptWrapper(iekeyname, enckey, NULL, NULL, FALSE);
}

EncryptWrapper::EncryptWrapper(const char *iekeyname, EncryptionKey *enckey, const char *iskeyname, SigningKey *sgnkey)
{
   InitEncryptWrapper(iekeyname, enckey, iskeyname, sgnkey, FALSE);
}

EncryptWrapper::EncryptWrapper(const char *iekeyname, EncryptionKey *enckey, int docompress)
{
   InitEncryptWrapper(iekeyname, enckey, NULL, NULL, docompress);
}

EncryptWrapper::EncryptWrapper(const char *iekeyname, EncryptionKey *enckey, const char *iskeyname, SigningKey *sgnkey, int docompress)
{
   InitEncryptWrapper(iekeyname, enckey, iskeyname, sgnkey, docompress);
}

EncryptWrapper::~EncryptWrapper()
{
   if (ekeyname) free(ekeyname);
   if (enckeycopy) delete enckeycopy;
   if (skeyname) free(skeyname);
   if (sgnkeycopy) delete sgnkeycopy;
   apr_pool_destroy(mp);
}

std::vector<char> *EncryptWrapper::EncryptBuffer(const char *buffer, size_t length)
{
   int smallblock = (ENV_READ_BUFFER > (length/2)) ? (length/2) : ENV_READ_BUFFER;
   size_t nlength = length;
   APRBucketBrigade *inputbb = new APRBucketBrigade(mp);
   if (inputbb->BufferToBucketBrigade(buffer, &length, smallblock) != 0)
   {
      errprint("failed to push the buffer into a bucket");
      delete inputbb;
      return NULL;
   }
   if (nlength != length)
   {
      errprint("bucket generation failed, pushed %d bytes to bucket, wanted %d", nlength, length);
      delete inputbb;
      return NULL;
   }
   APRBucketBrigade *outputbb = EncryptToBucket(inputbb);
   delete inputbb;
   char *outputbuffer = NULL;
   size_t ilength = 0;
   std::vector<char> *newstring = NULL;
   if (outputbb->BucketBrigadeToBuffer(&outputbuffer, &ilength) == 0)
   {
      // outputbuffer cannot be NULL!
      newstring = new std::vector<char>();
      newstring->insert(newstring->end(), outputbuffer, &(outputbuffer[ilength]));
   } else
   {
      errprint("failed to put bucket into buffer");
   }
   if (outputbuffer) free(outputbuffer);
   delete outputbb;
   return newstring;
}

APRBucketBrigade *EncryptWrapper::EncryptToBucket(APRBucketBrigade *bbrgde)
{
   CryptEChain *crypty = new CryptEChain();
   crypty->BuildEncryptionChain2(enckeycopy, sgnkeycopy, docompress, bbrgde->RealLength());
   APRBucketBrigade *outputbrgde = new APRBucketBrigade(mp);
   int prociter = 0;
   while (prociter == 0)
      prociter = crypty->ProcessIteration(bbrgde, outputbrgde);
   delete bbrgde;
   trusted = crypty->IsTrusted();
   delete crypty;
   if (prociter == -1)
   {
      delete outputbrgde;
      return NULL;
   }
   return outputbrgde;
}
